# 0. 导入必要的库
from util import load, get, dump # type: ignore
from lazypredict.Supervised import LazyClassifier
import os
import joblib

# 1. 加载训练集和测试集
X_train, X_test, y_train, y_test = load("X_train, X_test, y_train, y_test", f'{get("Xy_root")}/Xy')

# 2. 使用LazyClassifier进行快速模型评估
print("开始评估所有的模型:")
clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
models, predictions = clf.fit(X_train, X_test, y_train, y_test)

# 打印不同模型的评估结果对比
print(models)
# 检查models的列名
print("Columns in models DataFrame:", models.columns)
# 检查models的前几行数据
print(models.head())
# 提取模型名称和F1分数
scores = models[['F1 Score']].copy()
scores['Model'] = scores.index
print(scores)

# 3. 获取F1分数最高的模型
best_model_name = scores.loc[scores['F1 Score'].idxmax()]['Model']
print("F1分数最高的模型是: ", best_model_name)
# 根据模型名称，从模型字典中获取模型对象
best_model = clf.models[best_model_name]
# 确保目录存在
model_directory = get("model_root")
if not os.path.exists(model_directory):
    os.makedirs(model_directory)

# 4. 序列化最佳模型
model_path = f'{model_directory}/best_model.pkl'
joblib.dump(best_model, model_path)
print(f"最佳模型已保存到: {model_path}")
